package com.xfcode.ai.platform;

import cn.hutool.core.util.StrUtil;
import cn.hutool.http.ContentType;
import cn.hutool.json.JSONObject;
import com.google.gson.Gson;
import com.xfcode.ai.platform.bean.*;
import io.reactivex.BackpressureStrategy;
import io.reactivex.Flowable;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.jetbrains.annotations.NotNull;

import java.io.BufferedReader;
import java.io.IOException;
import java.util.Objects;

@Slf4j
public class MoonClient {
    private static final String DEFAULT_BASE_URL = "https://api.moonshot.cn/v1";
    private static final String CHAT_COMPLETION_SUFFIX = "/chat/completions";
    private static final String MODELS_SUFFIX = "/models";
    private static final String FILES_SUFFIX = "/files";
    private static final String ESTIMATE_TOKEN_COUNT = "/tokenizers/estimate-token-count";

    private String baseUrl;
    private String apiKey;
    private PlatformClient platformClient;

    public MoonClient(String apiKey) {
        this(apiKey, DEFAULT_BASE_URL);
        this.platformClient=new PlatformClient(apiKey,getFilesUrl(),getModelsUrl(),getEstimateTokenCountURL());
    }

    public PlatformClient getPlatformClient() {
        return platformClient;
    }

    public MoonClient(String apiKey, String baseUrl) {
        this.apiKey = apiKey;
        if (baseUrl.endsWith("/")) {
            baseUrl = baseUrl.substring(0, baseUrl.length() - 1);
        }
        this.baseUrl = baseUrl;
    }

    public String getChatCompletionUrl() {
        return baseUrl + CHAT_COMPLETION_SUFFIX;
    }
    public String getEstimateTokenCountURL() {
        return baseUrl + ESTIMATE_TOKEN_COUNT;
    }

    public String getModelsUrl() {
        return baseUrl + MODELS_SUFFIX;
    }

    public String getFilesUrl() {
        return baseUrl + FILES_SUFFIX;
    }

    public String getApiKey() {
        return apiKey;
    }

    public ModelsList getModels() throws IOException {
        okhttp3.OkHttpClient client = new okhttp3.OkHttpClient();
        okhttp3.Request request = new okhttp3.Request.Builder()
                .url(getModelsUrl())
                .addHeader("Authorization", "Bearer " + getApiKey())
                .build();
        try {
            okhttp3.Response response = client.newCall(request).execute();
            String body = response.body().string();
            Gson gson = new Gson();
            return gson.fromJson(body, ModelsList.class);
        } catch (IOException e) {
            e.printStackTrace();
            throw e;
        }
    }


    public ChatCompletionResponse chatCompletion(ChatCompletionRequest request) throws IOException {
        request.stream = false;
        okhttp3.OkHttpClient client = new okhttp3.OkHttpClient();
        okhttp3.MediaType mediaType = okhttp3.MediaType.parse("application/json");
        okhttp3.RequestBody body = okhttp3.RequestBody.create(mediaType, new Gson().toJson(request));
        okhttp3.Request httpRequest = new okhttp3.Request.Builder()
                .url(getChatCompletionUrl())
                .addHeader("Authorization", "Bearer " + getApiKey())
                .addHeader("Content-Type", "application/json")
                .post(body)
                .build();
        try {
            okhttp3.Response response = client.newCall(httpRequest).execute();
            String responseBody = response.body().string();
            Gson gson = new Gson();
            return gson.fromJson(responseBody, ChatCompletionResponse.class);
        } catch (IOException e) {
            e.printStackTrace();
            throw e;
        }
    }


    @SneakyThrows
    @Deprecated
    public  String chatCompletionDebug(ChatCompletionRequest request) {
        String requestBody = new JSONObject()
                .putOpt("model", request.getModel())
                .putOpt("messages", request.getMessages())
                .putOpt("stream", true)
                .toString();
        Request okhttpRequest = new Request.Builder()
                .url(getChatCompletionUrl())
                .post(RequestBody.create(MediaType.get(ContentType.JSON.getValue()), requestBody))
                .addHeader("Authorization", "Bearer " + apiKey)
                .build();
        Call call = new OkHttpClient().newCall(okhttpRequest);
        Response okhttpResponse = call.execute();
        BufferedReader reader = new BufferedReader(okhttpResponse.body().charStream());
        StringBuilder content = new StringBuilder();
        String line;
        while ((line = reader.readLine()) != null) {
            if (StrUtil.isBlank(line)) {
                continue;
            }
            content.append(line);
        }
        return  content.toString();
    }

    public String chatStreamCompletion(ChatCompletionRequest request) {
        StringBuffer sb =new StringBuffer();
        try {
                 chatCompletionStream(request).subscribe(
                    streamResponse -> {
                        if (streamResponse.getChoices().isEmpty()) {
                            return;
                        }
                        for (ChatCompletionStreamChoice choice : streamResponse.getChoices()) {
                            String finishReason = choice.getFinishReason();
                            if (finishReason != null) {
                                log.info("finish reason: " + finishReason);
                                continue;
                            }
                            sb.append(choice.getDelta().getContent());
                        }
                    },
                    error -> {
                        error.printStackTrace();
                    },
                    () -> {
                        log.info("complete");
                    }
            );
        } catch (Exception e) {
            throw new IllegalStateException("调用失败",e);
        }
        return  sb.toString();
    }

    // return a stream of ChatCompletionStreamResponse
    public Flowable<ChatCompletionStreamResponse> chatCompletionStream(ChatCompletionRequest request) throws IOException {
        request.stream = true;
        String content= new Gson().toJson(request);
        return getChatCompletionStreamResponseFlowable(content);
    }


    // return a stream of ChatCompletionStreamResponse
    public Flowable<ChatCompletionStreamResponse> chatCompletionStream(ChatCompletionFileRequest request) throws IOException {
        request.stream = true;
        String content= new Gson().toJson(request);
        return getChatCompletionStreamResponseFlowable(content);
    }


    private @NotNull Flowable<ChatCompletionStreamResponse> getChatCompletionStreamResponseFlowable( String content) throws IOException {
        OkHttpClient client = new OkHttpClient();
        MediaType mediaType = MediaType.parse("application/json");
        RequestBody body = RequestBody.create(mediaType,content);
        Request httpRequest = new Request.Builder()
                .url(getChatCompletionUrl())
                .addHeader("Authorization", "Bearer " + getApiKey())
                .addHeader("Content-Type", "application/json")
                .post(body)
                .build();
        Response response = client.newCall(httpRequest).execute();
        if (response.code() != 200) {
            throw new RuntimeException("Failed to start stream: " + response.body().string());
        }

        // get response body line by line
        return Flowable.create(emitter -> {
            ResponseBody responseBody = response.body();
            if (responseBody == null) {
                emitter.onError(new RuntimeException("Response body is null"));
                return;
            }
            String line;
            while ((line = responseBody.source().readUtf8Line()) != null) {
                if (line.startsWith("data:")) {
                    line = line.substring(5);
                    line = line.trim();
                }
                if (Objects.equals(line, "[DONE]")) {
                    emitter.onComplete();
                    return;
                }
                line = line.trim();
                if (line.isEmpty()) {
                    continue;
                }
                Gson gson = new Gson();
                ChatCompletionStreamResponse streamResponse = gson.fromJson(line, ChatCompletionStreamResponse.class);
                emitter.onNext(streamResponse);
            }
            emitter.onComplete();
        }, BackpressureStrategy.BUFFER);
    }
}
